import json
import platform
import re
import socket
import psutil
import os
import subprocess


def get_system_name():
    return platform.system()

def get_system_name_all():
    if get_system_name() == 'Windows':
        return platform.platform()
    else:
        return platform.linux_distribution()


def get_system_ip():
    return socket.gethostbyname(socket.gethostname())


def get_process_info(pid):
    try:
        # 根据进程ID获取进程对象
        process = psutil.Process(pid)
        # 获取进程信息
        pinfo = process.as_dict(attrs=['pid', 'name', 'cmdline', 'cpu_percent', 'memory_percent', 'status'])
        return pinfo
    except psutil.NoSuchProcess:
        # 如果没有找到对应的进程，返回 None
        return False


def restar_system():
    if get_system_name() == 'Windows':
        os.system("shutdown /r /t 0")
    else:
        os.system("reboot")


def set_firewall_rule(port, system, status, protocol, source_ip=None):
    if system == "Linux":
        if status == "open":
            # 在Linux下使用iptables开放端口
            if source_ip:
                subprocess.run(
                    ['iptables', '-A', 'INPUT', '-p', protocol, '--dport', str(port), '-s', source_ip, '-j', 'ACCEPT'])
            else:
                subprocess.run(['iptables', '-A', 'INPUT', '-p', protocol, '--dport', str(port), '-j', 'ACCEPT'])
            print(
                f"Opened port {port} on Linux firewall for {protocol} protocol{' from ' + source_ip if source_ip else ''}.")
        elif status == "close":
            # 在Linux下使用iptables关闭端口
            subprocess.run(['iptables', '-D', 'INPUT', '-p', protocol, '--dport', str(port), '-j', 'ACCEPT'])
            print(f"Closed port {port} on Linux firewall for {protocol} protocol.")
        else:
            print("Unsupported status. Please specify either 'open' or 'close'.")
    elif system == "Windows":
        if status == "open":
            # 在Windows下使用netsh开放端口
            if protocol == "TCP":
                os.system(
                    f'netsh advfirewall firewall add rule name="Open Port {port} TCP" dir=in action=allow protocol=TCP localport={port}{" remoteip=" + source_ip if source_ip else ""}')
                print(
                    f"Opened port {port} on Windows firewall for TCP protocol{' from ' + source_ip if source_ip else ''}.")
            elif protocol == "UDP":
                os.system(
                    f'netsh advfirewall firewall add rule name="Open Port {port} UDP" dir=in action=allow protocol=UDP localport={port}{" remoteip=" + source_ip if source_ip else ""}')
                print(
                    f"Opened port {port} on Windows firewall for UDP protocol{' from ' + source_ip if source_ip else ''}.")
            else:
                print("Unsupported protocol. Please specify either 'TCP' or 'UDP'.")
        elif status == "close":
            # 在Windows下使用netsh关闭端口
            rule_name = f"Open Port {port} {protocol}" if source_ip is None else f"Open Port {port} {protocol} from {source_ip}"
            os.system(f'netsh advfirewall firewall delete rule name="{rule_name}"')
            print(f"Closed port {port} on Windows firewall for {protocol} protocol.")
        else:
            print("Unsupported status. Please specify either 'open' or 'close'.")
    else:
        print("Unsupported system. Please specify either 'Linux' or 'Windows'.")


def get_firewall_rules(system):
    if system == "Linux":
        try:
            # 调用 iptables 命令获取端口规则信息
            output = subprocess.check_output('iptables -L INPUT -n --line-numbers', shell=True)

            # 将输出按行分割
            output_lines = output.decode('utf-8').splitlines()

            # 解析规则信息
            rules = []
            for line in output_lines[2:]:
                rule_info = line.split()
                if len(rule_info) >= 6:
                    rule = rule_info
                    rules.append(rule)

            # 将规则信息转换为 JSON 格式
            json_rules = json.dumps(rules, indent=4)
            return json_rules

        except Exception as e:
            return f"An error occurred: {e}"
    elif system == "Windows":

        try:
            # 调用 netsh 命令获取防火墙规则信息
            output = subprocess.check_output('netsh advfirewall firewall show rule name=all', shell=True)

            # 将输出按行分割
            output_lines = output.decode('gbk').splitlines()

            # 解析规则信息
            rules = []
            current_rule = {}
            for line in output_lines:
                if line.strip() == '':
                    if current_rule:
                        rules.append(current_rule)
                        current_rule = {}
                else:
                    key_value = line.split(':', 1)
                    if len(key_value) == 2:
                        key = key_value[0].strip()
                        value = key_value[1].strip()
                        current_rule[key] = value

            # 将规则信息转换为 JSON 格式
            json_rules = json.dumps(rules, ensure_ascii=False, indent=4)
            return json_rules

        except Exception as e:
            return f"An error occurred: {e}"
    else:
        print("Unsupported system. Please specify either 'Linux' or 'Windows'.")
        return

def off_system():
    if get_system_name() == 'Windows':
        os.system("shutdown /s /t 1")
    else:
        os.system("sudo shutdown now")
